package sample.methodremoveinstruct;

import org.apache.commons.lang3.StringUtils;
import org.objectweb.asm.*;
import util.FileUtils;

import java.io.IOException;

import static org.objectweb.asm.Opcodes.*;
import static util.StateMachineUtils.SEEN_NOTHING;

public class MethodRemovePrintStringDemo {

    public static void main(String[] args) throws IOException {
        Class<?> clazz = HelloWorld.class;
        byte[] bytes = invoke(clazz, Type.getInternalName(clazz) + "RemovePrintString.class");
        //  output
        FileUtils.writeBytes(Type.getInternalName(clazz) + "RemovePrintString.class", bytes);
        //  replace
        FileUtils.writeBytes(clazz.getResource("/").getPath() + Type.getInternalName(clazz) + ".class", bytes);
    }

    private static byte[] invoke(Class<?> clazz, String filePath) throws IOException {
        ClassReader classReader = new ClassReader(clazz.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_FRAMES);
        MethodRemovePrintStringVisitor visitor = new MethodRemovePrintStringVisitor(Opcodes.ASM9, classWriter);
        classReader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return classWriter.toByteArray();
    }

    protected static class MethodRemovePrintStringVisitor extends ClassVisitor {
        public MethodRemovePrintStringVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (mv != null && !"<init>".equals(name) && !"<clinit>".equals(name)) {
                boolean isAbstractMethod = (access & ACC_ABSTRACT) != 0;
                boolean isNativeMethod = (access & ACC_NATIVE) != 0;
                if (!isAbstractMethod && !isNativeMethod) {
                    mv = new MethodRemovePrintStringAdapter(api, mv);
                }
            }
            return mv;
        }

        private static class MethodRemovePrintStringAdapter extends MethodPatternAdapter {
            private final int SEEN_GETSTATIC = 1;
            private final int SEEN_GETSTATIC_LDC = 2;
            private String message;

            public MethodRemovePrintStringAdapter(int api, MethodVisitor mv) {
                super(api, mv);
            }

            @Override
            public void visitFieldInsn(int opcode, String owner, String name, String descriptor) {
                if (state == SEEN_NOTHING && opcode == GETSTATIC
                        && StringUtils.equals(owner, "java/lang/System")
                        && StringUtils.equals(name, "out")
                        && StringUtils.equals(descriptor, "Ljava/io/PrintStream;")) {
                    state = SEEN_GETSTATIC;
                    return;
                }
                super.visitFieldInsn(opcode, owner, name, descriptor);
            }

            @Override
            public void visitLdcInsn(Object value) {
                if (state == SEEN_GETSTATIC && value instanceof String) {
                    state = SEEN_GETSTATIC_LDC;
                    message = (String) value;
                    return;
                }
                super.visitLdcInsn(value);
            }

            @Override
            public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
                if (state == SEEN_GETSTATIC_LDC && opcode == INVOKEVIRTUAL
                        && StringUtils.equals(owner, "java/io/PrintStream")
                        && StringUtils.equals(name, "println")
                        && StringUtils.equals(descriptor, "(Ljava/lang/String;)V")) {
                    state = SEEN_NOTHING;
                    return;
                }
                super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
            }

            @Override
            protected void visitInsn() {
                switch (state) {
                    case SEEN_GETSTATIC:
                        mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                        break;
                    case SEEN_GETSTATIC_LDC:
                        mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                        mv.visitLdcInsn(message);
                        break;
                    default:
                        break;
                }
                state = SEEN_NOTHING;
            }
        }
    }
}
